import os
from langchain_community.llms import Tongyi
from langchain.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from dotenv import find_dotenv, load_dotenv

env = load_dotenv(find_dotenv())
# 这里需要将 DASHSCOPE_API_KEY 替换为你在阿里云控制台开通的 API KEY
os.environ["DASHSCOPE_API_KEY"] = os.getenv('ALI_API_KEY')

# 可以通过 model 指定模型
# llm = ChatTongyi(model='qwen-vl-max')
llm = Tongyi()

# 需要注意的一点是，这里需要指明具体的role，在这里是system和用户
prompt = ChatPromptTemplate.from_messages([
    ("system", "You are the technical writer"),
    ("user", "{input}")  # {input}为变量
])

output_parser = StrOutputParser()  # 输出string

# 我们可以把prompt和具体llm的调用和在一起（通过chain，chain可以理解为sequence of calls to take）
chain = prompt | llm | output_parser
res = chain.invoke({"input": "What is the Sora model?"})

print(res)
